import os

from setuptools import setup
from torch.utils import cpp_extension

# 获取当前目录
current_dir = os.path.dirname(os.path.abspath(__file__))

# 定义PyTorch自定义算子扩展
setup(
    name="dose_stat_torch",
    ext_modules=[
        cpp_extension.CppExtension(
            name="dose_stat_torch",
            sources=[os.path.join(current_dir, "DoseStatTorch.cpp")],
            extra_compile_args=["-std=c++14"],
        )
    ],
    cmdclass={"build_ext": cpp_extension.BuildExtension},
    packages=[""],
    package_dir={"": current_dir},
    package_data={
        "": ["*.py", "*.h"],
    },
    zip_safe=False,
)
